#include "inference.h"

int main(int argc, char* argv[])
{
    VideoCapture cap(VIDEO_PATH);
    Mat img;
    // -------- Step 1. Initialize OpenVINO Runtime Core --------
    ov::Core core;
    // -------- Step 2. Compile the Model --------
    auto compiled_model = core.compile_model(MODEL_PATH, "CPU");
    // -------- Step 3. Create an Inference Request --------
    ov::InferRequest infer_request = compiled_model.create_infer_request();
    TickMeter meter;
    // -------- Step 4.Read a picture file and do the preprocess --------
    while (true) {
        meter.start();
        cap.read(img);
        if (img.empty()) break;
        // Preprocess the image
        Mat letterbox_img = letterbox(img);
        float scale = letterbox_img.size[0] / 640; // 宽度
        Mat blob = blobFromImage(letterbox_img, 1.0 / 255.0, Size(640, 640), Scalar(), true); // 图像像素归一化
        // -------- Step 5. Feed the blob into the input node of the Model -------
        // Get input port for model with one input

        auto input_port = compiled_model.input();

        // Create tensor from external memory
        ov::Tensor input_tensor(input_port.get_element_type(), input_port.get_shape(), blob.ptr(0));
        std :: cout << "hello" << std :: endl;
        // Set input tensor for model with one input
        infer_request.set_input_tensor(input_tensor);
        // -------- Step 6. Start inference --------
        infer_request.infer();

        // -------- Step 7. Get the inference result --------
        auto output = infer_request.get_output_tensor(0);
        auto output_shape = output.get_shape();

        std :: cout << output_shape << std :: endl;
        // -------- Step 8. Postprocess the result --------
        float *data = output.data<float>();
        Mat output_buffer(output_shape[1], output_shape[2], CV_32F, data);
        transpose(output_buffer, output_buffer); //[8400,18]

        // 18: box[cx, cy, w, h] + Score + [5,2] keypoints
        std :: vector<yolo_Data> Created_Data;
        Created_Data= Create_Data(Created_Data, output_buffer, CONFIDENCE_THRESHOLD, scale);
        for (auto i : Created_Data) {
            //NMS
            std::vector<int> indices;
            NMSBoxes(i.boxes, i.class_scores, CONFIDENCE_THRESHOLD, NMS_THRESHOLD, indices);

            for (size_t j = 0; j < indices.size(); j++) {
                int index = indices[j]; // boxes,objects_keypoints中的索引
                drawBoundingBox(img, i.boxes[index], i.class_ids, i.class_scores[index], colors[i.class_ids]);
                draw_kpt(img, i.objects_keypoints[index], colors);
            }
        }
        namedWindow("YOLOv8-Pose OpenVINO Inference C++ Demo", WINDOW_AUTOSIZE);
        imshow("YOLOv8-Pose OpenVINO Inference C++ Demo", img);

        meter.stop();
        printf("Time: %f\t帧数：%f\n", meter.getTimeMilli(), 1000.0 / meter.getTimeMilli());
        meter.reset();

        if (waitKey(1) == 27) break;
    }

    destroyAllWindows();
    return 0;

}